#!/usr/bin/env python
# -*- encoding:utf8 -*-
# protoc-gen-erl
# Google's Protocol Buffers project, ported to Erlang.
# https://code.google.com/p/protoc-gen-erl/
#
# Copyright (c) 2010 , 林卓毅 (Zhuoyi Lin) netsnail@gmail.com
# All rights reserved.
#
# Use, modification and distribution are subject to the "New BSD License"
# as listed at <url: http://www.opensource.org/licenses/bsd-license.php >.

import sys, os
import os.path as path
from cStringIO import StringIO

#sys.path.append(os.path.dirname(os.path.realpath(__file__)))

import plugin_pb2
import google.protobuf.descriptor_pb2 as descriptor_pb2

_files = {}

def printerr(*args):
    sys.stderr.write(" ".join(args))
    sys.stderr.write("\n")
    sys.stderr.flush()

TYPE_PACKAGE = 1
TYPE_MESSAGE = 2

class TreeNode(object):
    def __init__(self, node_type, name = 'pb', parent=None):
        super(TreeNode, self).__init__()
        self.child = []
        self.parent = parent
        if parent:
            self.parent.add_child(self)
        self.name = name
        self.node_type = node_type 

    def add_child(self, child):
        self.child.append(child)

    def find_child(self, child_names):
        if child_names:
            for i in self.child:
                if i.name == child_names[0]:
                    return i.find_child(child_names[1:])
            raise KeyError(child_names)
        else:
            return self

    def get_child(self, child_name):
        for i in self.child:
            if i.name == child_name:
                return i
        return None

    def get_path(self):
        pos = self
        out = []
        while pos:
            out.append('.' + pos.name if pos.node_type in [TYPE_MESSAGE, TYPE_PACKAGE] else '')
            pos = pos.parent
        out.reverse()
        return ''.join(out)[1:]

    def get_global_name(self):
        return self.get_path()

    def get_local_name(self):
        pos = self
        out = []
        while pos and pos.node_type == TYPE_MESSAGE:
            out.append('.' + pos.name if pos.name else '')
            pos = pos.parent
        out.reverse()
        return ''.join(out)[1:]

    def get_package_name(self):
        pos = self
        while pos and pos.node_type == TYPE_MESSAGE:
            pos = pos.parent
        out = []
        while pos and pos.name != '':
            out.append(pos.name)
            pos = pos.parent
        out.reverse()
        return '.'.join(out)

    def __str__(self):
        return self.tostring(0)

    def __repr__(self):
        return str(self)
    
    def tostring(self, indent):
        if len(self.child) == 0:
            return ' ' * indent + '<TreeNode ' + self.name + ' ()>'
        else:
            return ' ' * indent + '<TreeNode ' + self.name + '(\n' +\
                    ',\n'.join([i.tostring(indent  + 4) for i in self.child]) + ')>'

    def get_root(self):
        pos = self
        while pos.parent:
            pos = pos.parent
        return pos

class PackageNode(TreeNode):
    def __init__(self, *arg, **kwarg):
        super(PackageNode, self).__init__(*arg, **kwarg)
        self.define = []
        self.record = []
        self.encode = []
        self.decode = []
        self.extend = []
        self.export = []
        self.includes = []
        self.types = []

    def gen(self, env):
        for child in self.child:
            if isinstance(child, PackageNode):
                child.gen(env)

        if not (self.define or self.record or self.encode or self.decode or self.extend):
            return

        hrl = Writer()
        write_header(hrl)
        package_name = self.get_package_name()
        erl_filename = package_name.replace('.', '/') 
        hrl_filename = package_name.replace('.', '_') 

        hrl("-ifndef(_%s_included).\n" % hrl_filename)
        hrl("-define(_%s_included, yeah).\n"% hrl_filename)
        for i in self.includes:
            node = env.refs[i]
            if node != self:
                hrl('-include("%s.hrl").\n' % node.get_global_name().replace('.', '_'))

        map(hrl, self.define)
        map(hrl, self.record)

        hrl("-endif.\n")

        erl = Writer()
        write_header(erl)
        erl("-module(%s).\n" % package_name)
        erl('-include("%s.hrl").\n' % hrl_filename)
        erl('-export([ types/0 ]).\n')
        erl('-export([ %s ]).\n' % ',\n'.join(self.export))
        erl('-import(protobuf).\n')
        erl('\n\n')
        map(erl, self.encode)
        map(erl, self.decode)
        map(erl, self.extend)
        erl('types() ->\n [%s].' % ', \n'.join(self.types))

        _files[erl_filename + '.erl'] = erl.getvalue()
        _files['include/' + hrl_filename + '.hrl'] = hrl.getvalue()

class Env(object):
    PROXY = set(['define', 'record', 'encode', 'decode', 'extend', 'export', 'includes', 'types'])

    def __init__(self):
        self.message_tree = PackageNode(TYPE_PACKAGE)
        self.scope = self.message_tree
        self.package = self.scope
        self.refs = {}

    def get_global_name(self):
        return self.scope.get_global_name()

    def get_local_name(self):
        return self.scope.get_local_name()

    def get_package_name(self):
        return self.scope.get_package_name()

    def lookup_name(self, name):
        if name[0] == '.':
            names = name.split('.')[1:]
            return self.message_tree.find_child(names)
        else:
            names = name.split('.')
            try:
                return self.scope.find_child(names) 
            except:
                return self.package.find_child(names)

    def enter_package(self, package):
        pos = self.message_tree
        if package != '':
            names = package.split('.')
            for i, name in enumerate(names):
                new_pos = pos.get_child(name)
                if new_pos:
                    pos = new_pos
                else:
                    return self._build_nodes(pos, names[i:])
        return pos

    def enter_file(self, filename, package):
        self.scope = self.enter_package(package)
        self.package = self.scope
        self.refs[filename] = self.package 

    def exit_file(self):
        self.scope = self.scope.parent
        self.package = self.scope

    def enter(self, message_name):
        try:
            node = self.lookup_name(message_name)
            self.scope = node
        except KeyError:
            self.scope = TreeNode(TYPE_MESSAGE, message_name, self.scope)

    def exit(self):
        self.scope = self.scope.parent

    def _build_nodes(self, node, names):
        parent = node
        for i in names:
            parent = PackageNode(TYPE_PACKAGE, i, parent)
        return parent

    def __getattr__(self, key):
        if key in self.PROXY:
            return getattr(self.package, key)
        raise StandardError

class Writer(object):
    def __init__(self):
        self.io = StringIO()
        self.__indent = ''

    def getvalue(self):
        return self.io.getvalue()

    def __enter__(self):
        self.__indent += '    '
        return self

    def __exit__(self, type, value, trackback):
        self.__indent = self.__indent[:-4]

    def __call__(self, data):
        self.io.write(self.__indent)
        self.io.write(data)

def code_gen_enum(enum_desc, env, is_gen):
    if not is_gen:
        return
    template = "-define('%s', %d).\n"
    for enum_value in enum_desc.value:
        env.enter(enum_value.name)    
        define = template%(
                env.get_global_name(),
                enum_value.number)
        env.define.append(define)
        env.exit()

FDP = plugin_pb2.descriptor_pb2.FieldDescriptorProto
_field_dct = {
    FDP.TYPE_DOUBLE:"double",
    FDP.TYPE_FLOAT:"float",
    FDP.TYPE_INT64:"int64",
    FDP.TYPE_UINT64:"uint64",
    FDP.TYPE_INT32:"int32",
    FDP.TYPE_FIXED64:"fixed64",
    FDP.TYPE_FIXED32:"fixed32",
    FDP.TYPE_BOOL:"bool",
    FDP.TYPE_STRING:"string",
    #FDP.TYPE_GROUP:"",
    #FDP.TYPE_MESSAGE:"",
    FDP.TYPE_BYTES:"bytes",
    FDP.TYPE_UINT32:"uint32",
    FDP.TYPE_ENUM:"enum",
    FDP.TYPE_SFIXED32:"sfixed32",
    FDP.TYPE_SFIXED64:"sfixed64",
    FDP.TYPE_SINT32:"sint32",
    FDP.TYPE_SINT64:"sint64",
}

def get_dec_func(type_name, env):
    node = env.lookup_name(type_name)
    module = node.get_package_name() 
    local_name = node.get_local_name()
    if module == env.get_package_name():
        return "fun 'dec_%s'/2"%local_name
    else:
        return "fun '%s':'dec_%s'/2"%(module, local_name)

def get_enc_func(type_name, env):
    node = env.lookup_name(type_name)
    module = node.get_package_name() 
    local_name = node.get_local_name()
    if module == env.get_package_name():
        return "fun 'enc_%s'/2"%local_name
    else:
        return "fun '%s':'enc_%s'/2"%(module, local_name)

def get_default_value(env, field_desc):
        value = field_desc.default_value
        if field_desc.type == FDP.TYPE_STRING:
            value = '"' + value + '"'
        elif field_desc.type == FDP.TYPE_ENUM:
            node = env.lookup_name(value)
            value = "?'" + node.get_global_name() + "'"
        return value

def make_decoder_type(field_desc, env):
    return """{%d, %s, #'%s'.%s, %s, %s, <<"%s">>, %s}"""%(
        field_desc.number,
        'bytes' if field_desc.type == FDP.TYPE_MESSAGE else _field_dct[field_desc.type],
        env.get_global_name(),
        field_desc.name,
        get_dec_func(field_desc.type_name, env) \
            if field_desc.type == FDP.TYPE_MESSAGE else "fun protobuf:nop/2",
        'true' if field_desc.label == field_desc.LABEL_REPEATED else 'false',
        field_desc.name,
        get_default_value(env, field_desc) if field_desc.HasField('default_value') else 'undefined'
        )

def make_extend_decoder_type(field_desc, extension_name, env):
    node = env.lookup_name(field_desc.extendee)
    extendee = node.get_global_name()

    func = get_dec_func(field_desc.type_name, env)
    return """{'%s', %d, %s, #'%s'.%s, %s, %s, <<"%s">>, %s}"""%(
        extendee,
        field_desc.number,
        'bytes' if field_desc.type == FDP.TYPE_MESSAGE else _field_dct[field_desc.type],
        extendee,
        extension_name,
        func if field_desc.type == FDP.TYPE_MESSAGE else "fun protobuf:nop/2",
        'true' if field_desc.label == field_desc.LABEL_REPEATED else 'false',
        field_desc.name,
        get_default_value(env, field_desc) if field_desc.HasField('default_value') else 'undefined'
        )

def make_extend_encoder_type(field_desc, extension_name, env):
    node = env.lookup_name(field_desc.extendee)
    extendee = node.get_global_name()
    
    func = get_enc_func(field_desc.type_name, env)
    if field_desc.type == FDP.TYPE_MESSAGE:
        return "{'%s', %d, %s, true, %s, '%s'}"%(
            extendee,
            field_desc.number,
            func,
            'true' if field_desc.label == field_desc.LABEL_REPEATED else 'false',
            env.get_global_name()
            )
    else:
        return "{'%s', %d, true, %s, %s, '%s'}"%(
            extendee,
            field_desc.number,
            _field_dct[field_desc.type],
            'true' if field_desc.label == field_desc.LABEL_REPEATED else 'false', 
            env.get_global_name()
            )

def code_gen_extensions(field_desc, local_name, env, is_gen):
    if not is_gen:
        return 
    extend = Writer()
    env.export.append("'extend_dec_%s'/0"%local_name)
    extend("'extend_dec_%s'()->\n"%local_name)
    with extend:
        extend(
            make_extend_decoder_type(field_desc, 'extensions', env)
            + ".\n\n")
    env.extend.append(extend.getvalue())

    extend = Writer()
    env.export.append("'extend_enc_%s'/0"%local_name)
    extend("'extend_enc_%s'()->\n"%local_name)
    with extend:
        extend(
            make_extend_encoder_type(field_desc, 'extensions', env)
            + ".\n\n")
    env.extend.append(extend.getvalue())

def code_gen_message(message_descriptor, env, is_gen):
    env.enter(message_descriptor.name)
    global_message_name = env.get_global_name()
    local_message_name = env.get_local_name()

    for enum_desc in message_descriptor.enum_type:
        code_gen_enum(enum_desc, env, is_gen)

    for msg_desc in message_descriptor.nested_type:
        code_gen_message(msg_desc, env, is_gen)

    if not is_gen:
        env.exit()
        return 

    fields = []
    for i, field_desc in enumerate(message_descriptor.field):
        if field_desc.HasField('default_value'):
            value = "%s = %s"%(field_desc.name, get_default_value(env, field_desc))
        else:
            value = field_desc.name
        fields.append(value)

    has_extensions = len(message_descriptor.extension_range) > 0
    has_extend = len(message_descriptor.extension) > 0

    encoder = Writer()
    encoder("'enc_%s'(Data) when is_record(Data, '%s')->\n" % (
        local_message_name, global_message_name))
    with encoder:
        encoder("'enc_%s'(Data, []).\n\n" % local_message_name)
    encoder("'enc_%s'(undefined, _Extensions) ->\n" %local_message_name)
    with encoder:
        encoder('<<>>;\n')
    extension_used = False
    body = Writer()
    with body:
        body("iolist_to_binary([\n")
        with body:
            for i, field_desc in enumerate(message_descriptor.field):
                    if i + 1 < len(message_descriptor.field) or has_extensions:
                        end = ',\n'
                    else:
                        end = '\n'
                    if field_desc.label == field_desc.LABEL_REPEATED:
                        if field_desc.type == FDP.TYPE_MESSAGE:
                            extension_used = True
                            body(
                                "[protobuf:encode(%d, %s(I, Extensions), %s) || I <- Data#'%s'.%s ]%s" % (
                                    field_desc.number,
                                    get_enc_func(field_desc.type_name, env),
                                    "bytes",
                                    global_message_name,
                                    field_desc.name,
                                    end))
                        else:
                            field_type = _field_dct[field_desc.type]
                            body(
                                "[protobuf:encode(%d, I, %s) || I <- Data#'%s'.%s ]%s" % (
                                    field_desc.number,
                                    field_type,
                                    global_message_name,
                                    field_desc.name,
                                    end))
                    else:
                        if field_desc.type == FDP.TYPE_MESSAGE:
                            extension_used = True
                            body("protobuf:encode(%d, %s(Data#'%s'.%s, Extensions), %s)%s" % (
                                field_desc.number,
                                get_enc_func(field_desc.type_name, env),
                                global_message_name,
                                field_desc.name,
                                "bytes",
                                end))
                        else:
                            field_type = _field_dct[field_desc.type]
                            body("protobuf:encode(%d, Data#'%s'.%s, %s)%s" % (
                                field_desc.number, global_message_name, field_desc.name, field_type, end))
            if has_extensions:
                extension_used = True
                body(
                    "protobuf:encode_extensions('%s', Data#'%s'.extensions, Extensions)\n"%(
                        global_message_name, global_message_name
                    ))
        body("]).\n\n")
    if extension_used == True:
        encoder("'enc_%s'(Data, Extensions) when is_record(Data, '%s')->\n" % (
            local_message_name, global_message_name))
    else:
        encoder("'enc_%s'(Data, _Extensions) when is_record(Data, '%s')->\n" % (
            local_message_name, global_message_name))

    encoder(body.getvalue())

    decoder = Writer()
    types = "'dec_types_%s'" % local_message_name
    decoder("%s()->\n" % types)  
    with decoder:
        decoder("[\n")
        with decoder:
            for i, field_desc in enumerate(message_descriptor.field):
                if i + 1 < len(message_descriptor.field):
                    end = ',\n'
                else:
                    end = '\n'
                decoder(
                    make_decoder_type(field_desc, env)
                    + end)
        decoder("].\n")

    decoder("'dec_%s'(Data) ->\n" % local_message_name)
    with decoder:
        decoder("'dec_%s'(Data, []).\n\n" % local_message_name)
    decoder("'dec_%s'(Data, Extensions) ->\n" % local_message_name)
    with decoder:
        decoder("protobuf:decode_loop(Data, #'%s'{}, %s(), Extensions).\n\n" \
                %(global_message_name, types))

    for field_desc in message_descriptor.extension:
        code_gen_extensions(field_desc,
                            local_message_name + '.' + field_desc.name,
                            env, is_gen)

    if len(fields):
        fields = ', '.join(fields) 
    else:
        fields = ''
    record = "-record('%s', {%s%s}). \n"%(
        global_message_name, fields,
        ', extensions' if has_extensions else '')
    env.record.append(record)
    env.encode.append(encoder.getvalue())
    env.decode.append(decoder.getvalue())
    env.export.append("'enc_%s'/1" % local_message_name)
    env.export.append("'enc_%s'/2" % local_message_name)
    env.export.append("'dec_%s'/1" % local_message_name)
    env.export.append("'dec_%s'/2" % local_message_name)
    env.export.append("'dec_types_%s'/0" % local_message_name)

    env.types.append(
            "{'%(global)s', 'enc_%(local)s', 'dec_%(local)s'}" % \
                    {'global': global_message_name, 'local' : local_message_name})
    env.exit()

def write_header(writer):
    writer("""%% Generated By protoc-gen-erl Do not Edit
""")

def code_gen_file(proto_file, env, is_gen):
    filename = path.splitext(proto_file.name)[0]
    env.enter_file(filename, proto_file.package)

    for f in proto_file.dependency:
        inc_file = path.splitext(f)[0]
        env.includes.append(inc_file)

    for enum_desc in proto_file.enum_type:
        code_gen_enum(enum_desc, env, is_gen)

    for msg_desc in proto_file.message_type:
        code_gen_message(msg_desc, env, False)

    for msg_desc in proto_file.message_type:
        code_gen_message(msg_desc, env, is_gen)

    for field_desc in proto_file.extension:
        code_gen_extensions(field_desc, field_desc.name, env, is_gen)

    env.exit_file()

    
def main():
    plugin_require_bin = sys.stdin.read()
    code_gen_req = plugin_pb2.CodeGeneratorRequest()
    code_gen_req.ParseFromString(plugin_require_bin)

    env = Env()
    for proto_file in code_gen_req.proto_file:
        code_gen_file(proto_file, env,
                proto_file.name in code_gen_req.file_to_generate)

    code_generated = plugin_pb2.CodeGeneratorResponse()
    
    env.message_tree.gen(env)
    for k in  _files:
        file_desc = code_generated.file.add()
        file_desc.name = k
        file_desc.content = _files[k]

    sys.stdout.write(code_generated.SerializePartialToString())

if __name__ == "__main__":
    main()


